import os
import random
from sklearn.model_selection import train_test_split

# 数据集根目录
data_root = "/home/aistudio/data/data89094/Garbages"
# 各类别文件夹
classes = ["Harmful", "Kitchen", "Other", "Recyclable"]

# 生成 labels.txt：列出类别，每行一个类别，行号为类别 ID（从 0 开始）
with open(os.path.join(data_root, "labels.txt"), "w", encoding="utf-8") as f:
    for cls in classes:
        f.write(cls + "\n")

# 收集所有图片路径和对应标签
image_paths = []
labels = []
for label_id, cls in enumerate(classes):
    cls_dir = os.path.join(data_root, cls)
    for img_name in os.listdir(cls_dir):
        img_path = os.path.join(cls_dir, img_name)
        image_paths.append(img_path)
        labels.append(label_id)

# 划分训练集和测试集，测试集占比可自行调整，这里设为 0.2（即 20% 作为测试集）
train_paths, test_paths, train_labels, test_labels = train_test_split(
    image_paths, labels, test_size=0.2, random_state=42
)

# 生成 train_list.txt：每行格式为 "图片路径 标签"
with open(os.path.join(data_root, "train_list.txt"), "w", encoding="utf-8") as f:
    for path, label in zip(train_paths, train_labels):
        f.write(f"{path} {label}\n")

# 生成 test_list.txt：每行格式为 "图片路径 标签"
with open(os.path.join(data_root, "test_list.txt"), "w", encoding="utf-8") as f:
    for path, label in zip(test_paths, test_labels):
        f.write(f"{path} {label}\n")